import math

from core.cuda import cuda_module
from core.optimizer import Optimizer


class Adam(Optimizer):
    def __init__(self, alpha=0.001, beta1=0.9, beta2=0.999, eps=1e-8):
        """
        Adam 优化器，结合了动量和 RMSProp 策略。

        Args:
            alpha (float): 学习率，控制参数更新的步长。
            beta1 (float): 一阶矩估计的衰减系数，通常为 0.9。
            beta2 (float): 二阶矩估计的衰减系数，通常为 0.999。
            eps (float): 为了数值稳定性而添加的小常数，一般取 1e-8。
        """
        super().__init__()
        self.t = 0
        self.alpha = alpha
        self.beta1 = beta1
        self.beta2 = beta2
        self.eps = eps
        self.ms = {}  # 一阶矩估计
        self.vs = {}  # 二阶矩估计

    def update(self, *args, **kwargs):
        """
        更新优化器状态，记录迭代次数。

        Args:
            *args: 任意位置参数。
            **kwargs: 任意关键字参数。

        Returns:
            None
        """
        self.t += 1
        super().update(*args, **kwargs)

    @property
    def lr(self):
        """
        计算当前的学习率。

        Returns:
            float: 当前学习率。
        """
        fix1 = 1. - math.pow(self.beta1, self.t)
        fix2 = 1. - math.pow(self.beta2, self.t)
        return self.alpha * math.sqrt(fix2) / fix1

    def update_one(self, param):
        """
        更新单个参数。

        Args:
            param (Parameter): 待更新的参数。

        Returns:
            None
        """
        xp = cuda_module

        key = id(param)
        if key not in self.ms:
            self.ms[key] = xp.zeros_like(param.data)
            self.vs[key] = xp.zeros_like(param.data)

        m, v = self.ms[key], self.vs[key]
        beta1, beta2, eps = self.beta1, self.beta2, self.eps
        grad = param.grad.data

        m += (1 - beta1) * (grad - m)
        v += (1 - beta2) * (grad * grad - v)

        param.data = param.data - self.lr * m / (xp.sqrt(v) + eps)
